#-------------------------------------------------------------------------------
# Name:        taskfile
# Purpose:
#
# Author:      SUPER_紫电  QQ:911344755
#
# Created:     07/09/2022
# Copyright:   (c) SUPER_紫电 2022 All rights reserved.
# Licence:     <Modified BSD License>
#-------------------------------------------------------------------------------

use_dll = True

import os
import sys
try:
    sys.path.append(os.path.dirname(__file__) + r'\..\OpenLibSys')
except:
    sys.path.append(os.getcwd() + r'\..\OpenLibSys')

import struct
from time import time, sleep
from ctypes import addressof, byref, cast, c_buffer, memmove, memset, sizeof
import pci as pci
from ata import *
if use_dll:
    import OpenLibDll as ols
else:
    import OpenLibSys as ols

BIT0 = AMNF = ERR  = 0x01
BIT1 = TONF = INDX = 0x02
BIT2 = ABRT = CORR = 0x04
BIT3 = MCR  = DRQ  = 0x08
BIT4 = IDNF = DSC  = 0x10
BIT5 = MC   = DF   = 0x20
BIT6 = UNC  = DRDY = 0x40
BIT7 = BBK  = BSY  = 0x80

ata_controller = []
ata_port = {
    'base': 0x0,
    'alt' : 0x0,
}

ata_reg = {
    'data': {'address': 0x0, 'value': 0x0}, # Data register
    'fe'  : {'address': 0x0, 'value': 0x0}, # 1 - Features  | Error register
    'sc'  : {'address': 0x0, 'value': 0x0}, # 2 - Sector Count | Interrupt Reason register
    'low' : {'address': 0x0, 'value': 0x0}, # 3 - LBA Low register
    'mid' : {'address': 0x0, 'value': 0x0}, # 4 - LBA Mid/Byte Count Low register
    'high': {'address': 0x0, 'value': 0x0}, # 5 - LBA High/Byte Count High register
    'dev' : {'address': 0x0, 'value': 0x0}, # 6 - Device register
    'cmd' : {'address': 0x0, 'value': 0x0}, # 7 - Command | Status register
    'alt' : {'address': 0x0, 'value': 0x0}, # Device Control | Alternate Status register
}

drv_initialized = False
drv_hmodule = None
cur_error  = 0xEE
cur_status = 0xEE
ident_dev_data = IDENTIFY_DEVICE_DATA()

def inp(port): return ols.ReadIoPortByte(port)
def inpw(port): return ols.ReadIoPortWord(port)
def inpd(port): return ols.ReadIoPortDword(port)
def outp(port, value): return ols.WriteIoPortByte(port, value)
def outpw(port, value): return ols.WriteIoPortWord(port, value)
def outpd(port, value): return ols.WriteIoPortDword(port, value)

def init_drv():
    global drv_initialized
    global drv_hmodule
    if use_dll:
        drv_initialized = ols.InitOpenLibSys(drv_hmodule)
    else:
        drv_initialized = ols.InitOpenLibSys()
    return drv_initialized

def deinit_drv():
    if drv_initialized:
        if use_dll:
            ols.DeinitOpenLibSys(drv_hmodule)
        else:
            ols.DeinitOpenLibSys()

def enum_ata_controller(ata_controller):
    bar = pci.PCICFG_REG_BAR0
    for bus in range(256):
        for dev in range(32):
            for func in range(8):
                address = pci.PciBusDevFunc(bus, dev, func)
                vidValue = ols.ReadPciConfigDword(address, 0x00)
                if (vidValue == 0xFFFFFFFF) or (vidValue == 0):
                    continue
                classValue = ols.ReadPciConfigDword(address, 0x08)
                headerType = ols.ReadPciConfigByte(address, 0x0E)
                if classValue >> 24 & 0xFF != 1:
                    continue
                if headerType != 0:
                    continue
                for i in range(3):
                    base = ols.ReadPciConfigDword(address, bar + i * 8)
                    alt = ols.ReadPciConfigDword(address, bar + i * 8 + 4)
                    if(base&1) and (alt&1) and (base>0x170) and (alt<0xFFFF):
                        ata_port = {}
                        ata_port['base'] = base - 1
                        ata_port['alt'] = alt + 1
                        ata_controller.append(ata_port)

def set_ata_controller():
    global ata_controller
    ata_controller.clear()
    enum_ata_controller(ata_controller)

def device_select(device):
    global ata_reg
    if device == 0:
        ata_reg['dev']['value'] = 0x40
    elif device == 1:
        ata_reg['dev']['value'] = 0x50

def set_reg_addr(base, alt, dev=0):
    global ata_port
    global ata_reg
    ata_port['base'] = base
    ata_port['alt'] = alt
    device_select(dev)
    ata_reg['data']['address'] = base
    ata_reg['fe']['address']   = base + 1
    ata_reg['sc']['address']   = base + 2
    ata_reg['low']['address']  = base + 3
    ata_reg['mid']['address']  = base + 4
    ata_reg['high']['address'] = base + 5
    ata_reg['dev']['address']  = base + 6
    ata_reg['cmd']['address']  = base + 7
    ata_reg['alt']['address']  = alt

def get_reg_value():
    global ata_reg
    global cur_error
    global cur_status
    ata_reg['fe']['value']   = inp(ata_reg['fe']['address'])
    ata_reg['sc']['value']   = inp(ata_reg['sc']['address'])
    ata_reg['low']['value']  = inp(ata_reg['low']['address'])
    ata_reg['mid']['value']  = inp(ata_reg['mid']['address'])
    ata_reg['high']['value'] = inp(ata_reg['high']['address'])
    ata_reg['dev']['value']  = inp(ata_reg['dev']['address'])
    ata_reg['cmd']['value']  = inp(ata_reg['cmd']['address'])
    ata_reg['alt']['value']  = inp(ata_reg['alt']['address'])
    cur_error = ata_reg['fe']['value']
    cur_status = ata_reg['alt']['value']

def reg_error(bit):
    return inp(ata_reg['fe']['address']) & bit

def reg_status(bit):
    return inp(ata_reg['alt']['address']) & bit

def eob():
    return reg_status(ERR) or reg_status(BSY)

def wait_not_busy(timeout=10):
    result = True
    if reg_status(BSY):
        start = time()
        while reg_status(BSY):
            if time() - start > timeout:
                result = False
                break
    return result

def eto(timeout):
    return reg_status(ERR) or not wait_not_busy(timeout)

def write_lba(ulba, sc, cmd, fe=0):
    dev = ata_reg['dev']['value']
    lba = bytearray(6)
    for i in range(6):
        lba[i] = ulba >> 8
        ulba >>= 8
    outp(ata_reg['fe']['address'], fe >> 8)
    outp(ata_reg['sc']['address'], sc >> 8)
    outp(ata_reg['low']['address'], lba[3])
    outp(ata_reg['mid']['address'], lba[4])
    outp(ata_reg['high']['address'], lba[5])
    outp(ata_reg['fe']['address'], fe & 0xFF)
    outp(ata_reg['sc']['address'], sc & 0xFF)
    outp(ata_reg['low']['address'], lba[0])
    outp(ata_reg['mid']['address'], lba[1])
    outp(ata_reg['high']['address'], lba[2])
    outp(ata_reg['dev']['address'], dev)
    outp(ata_reg['cmd']['address'], cmd)

def lba_to_bytes(ulba):
    lba = bytearray(6)
    for i in range(6, -1, -1):
        lba[i] = ulba & 0xFF
        ulba >>= 8
    return lba

def get_lba():
    lba = bytearray(6)
    lba[3] = inp(ata_reg['high']['address'])
    lba[4] = inp(ata_reg['mid']['address'])
    lba[5] = inp(ata_reg['low']['address'])
    outp(ata_reg['alt']['address'], 0x80)
    lba[0] = inp(ata_reg['high']['address'])
    lba[1] = inp(ata_reg['mid']['address'])
    lba[2] = inp(ata_reg['low']['address'])
    ulba = 0
    for i in range(5):
        ulba = ulba + lba[i]
        ulba <<= 8
    return ulba + lba[5]

def get_max_lba():
    if not ata_cmd(cmd=0x27):
        return 0
    return get_lba()

def set_max_lba(ulba):
    if ulba > get_max_lba():
        return False
    lba = lba_to_bytes(ulba)
    dev = ata_reg['dev']['value']
    outp(ata_reg['low']['address'], lba[2])
    outp(ata_reg['low']['address'], lba[5])
    outp(ata_reg['mid']['address'], lba[1])
    outp(ata_reg['mid']['address'], lba[4])
    outp(ata_reg['high']['address'], lba[0])
    outp(ata_reg['high']['address'], lba[3])
    outp(ata_reg['dev']['address'], dev)
    outp(ata_reg['cmd']['address'], 0x37)
    return ulba == get_lba()

def read_verify_lba(ulba, sc, timeout=10):
    write_lba(ulba, sc, READ_VERIFY_SECTORS_EXT)
    start = time()
    while True:
        while reg_status(BSY):
            if time() - start > timeout:
                cmd_recal()
                return -6
        if reg_status(INDX):
            cmd_recal()
            ulba -= 1
            return -15
        stat = inp(ata_reg['alt']['address'])
        if stat == 0x10:
            return -16
        if stat == 0x50 or stat == 0x51:
            break
    if reg_error(AMNF):
        cmd_recal()
        return -7
    if reg_status(ERR):
        result = -8 if reg_error(TONF) else -1
        result = -9 if reg_error(ABRT) else -1
        result = -10 if reg_error(IDNF) else -1
        result = -11 if reg_error(UNC) else -1
        result = -12 if reg_error(BBK) else -1
        ulba = get_lba()
        cmd_recal()
    else:
        result = time() - start
        ulba = ulba + sc - 1
    return result

def write_uncorrectable(ulba, sc):
    write_lba(ulba, sc, WRITE_UNCORRECTABLE_EXT, 0xA5)

def send_ata_cmd(fe, sc, low, mid, high, dev, cmd, timeout=10):
    stat = inp(ata_reg['alt']['address'])
    if (stat & ERR) or (stat & BSY) or (stat & DRQ):
        return False
    outp(ata_reg['fe']['address'], fe >> 8)
    outp(ata_reg['sc']['address'], sc >> 8)
    outp(ata_reg['low']['address'], low >> 8)
    outp(ata_reg['mid']['address'], mid >> 8)
    outp(ata_reg['high']['address'], high >> 8)
    outp(ata_reg['fe']['address'], fe & 0xFF)
    outp(ata_reg['sc']['address'], sc & 0xFF)
    outp(ata_reg['low']['address'], low & 0xFF)
    outp(ata_reg['mid']['address'], mid & 0xFF)
    outp(ata_reg['high']['address'], high & 0xFF)
    outp(ata_reg['dev']['address'], dev)
    outp(ata_reg['cmd']['address'], cmd)
    if not wait_not_busy(timeout):
        return False
    stat = inp(ata_reg['alt']['address'])
    return ~(stat & ERR) and (stat != 0)

def ata_cmd(fe=0, sc=0, low=0, mid=0, high=0, cmd=0, timeout=10):
    dev = ata_reg['dev']['value']
    return send_ata_cmd(fe, sc, low, mid, high, dev, cmd, timeout)

def soft_reset():
    outp(ata_reg['alt']['address'], 4)
    sleep(0.001)
    outp(ata_reg['alt']['address'], 0)

def cmd_identify():
    return ata_cmd(cmd=0xEC) and reg_status(DRQ)

def cmd_sleep_drive():
    return ata_cmd(cmd=0xE6)

def cmd_standby():
    return ata_cmd(cmd=0xE0, timeout=3)

def cmd_idle():
    return ata_cmd(cmd=0xE1, timeout=30)

def cmd_wake_drive():
    soft_reset()
    outp(ata_reg['cmd']['address'], 0xE0)
    if not wait_not_busy(3):
        return False
    outp(ata_reg['cmd']['address'], 0xE1)
    return wait_not_busy(30)

def cmd_recal():
    return ata_cmd(cmd=0x10, timeout=30)

def cmd_rdy():
    return send_ata_cmd(0,0,0,0,0,0x10,0)

def cmd_set_pwd():
    return ata_cmd(cmd=0xF1) and reg_status(DRQ)

def cmd_unlock_unit():
    return ata_cmd(cmd=0xF2) and reg_status(DRQ)

def cmd_erase_prep():
    return ata_cmd(cmd=0xF3)

def cmd_erase_unit():
    return ata_cmd(cmd=0xF4) and reg_status(DRQ)

def cmd_freeze_lock():
    return ata_cmd(cmd=0xF5)

def cmd_disable_pwd():
    return ata_cmd(cmd=0xF6) and reg_status(DRQ)

def cmd_smart_on():
    return ata_cmd(0xD8,0,0,0x4F,0xC2,0xB0) and reg_status(DRDY)

def read_file(filename, buffer):
    try:
        with open(filename, 'rb') as file:
            while True:
                data = file.read(4)
                if not data:
                    break
                buffer.extend(struct.unpack('@4B', data))
        return True
    except:
        return False

def save_file(filename, buffer):
    try:
        with open(filename, 'wb') as file:
            file.write(buffer)
        return True
    except:
        return False

"""
def read_data(buffer, offset=0, count=0):
    if count == 0:
        count = len(buffer) // 512
    while count > 0:
        stat = inp(ata_reg['alt']['address'])
        while (stat & DRQ) or (stat & BSY) and ~(stat & ERR):
            if (stat & BSY):
                start = time()
                while (stat & BSY):
                    if time() - start > 10:
                        if offset == len(buffer):
                            break
                        else:
                            return False
                    stat = inp(ata_reg['alt']['address'])
            if (stat & DRQ):
                try:
                    data = inpd(ata_reg['data']['address'])
                    buffer[offset+0] = data & 0xFF
                    buffer[offset+1] = data >> 8 & 0xFF
                    buffer[offset+2] = data >> 16 & 0xFF
                    buffer[offset+3] = data >> 24
                    offset += 4
                except:
                    return False
            stat = inp(ata_reg['alt']['address'])
            if (stat & ERR):
                return False
        count -= 256
    return count <= 0

def write_data(buffer, offset=0, count=0):
    if count == 0:
        count = len(buffer) // 512
    while count > 0:
        stat = inp(ata_reg['alt']['address'])
        while (stat & DRQ) or (stat & BSY) and ~(stat & ERR):
            if (stat & BSY):
                start = time()
                while (stat & BSY):
                    if time() - start > 3:
                        if offset == len(buffer):
                            break
                        else:
                            return False
                    stat = inp(ata_reg['alt']['address'])
            if (stat & DRQ):
                try:
                    if is_64bit_os():
                        data = struct.unpack_from('@L', buffer, offset)[0]
                        outpd(ata_reg['data']['address'], data)
                        offset += 4
                    else:
                        data = struct.unpack_from('@H', buffer, offset)[0]
                        outpw(ata_reg['data']['address'], data)
                        offset += 2
                except:
                    return False
            stat = inp(ata_reg['alt']['address'])
            if (stat & ERR):
                return False
        count -= 256
    return count <= 0
"""

def read_data(buffer):
    count = len(buffer) // 512
    if count < 0:
        return False
    offset = 0
    stat = inp(ata_reg['alt']['address'])
    while (stat & DRQ) or (stat & BSY) and ~(stat & ERR):
        if (stat & BSY):
            while (stat & BSY):
                stat = inp(ata_reg['alt']['address'])
        if (stat & DRQ):
            try:
                data = inpd(ata_reg['data']['address'])
                buffer[offset+0] = data & 0xFF
                buffer[offset+1] = data >> 8 & 0xFF
                buffer[offset+2] = data >> 16 & 0xFF
                buffer[offset+3] = data >> 24
                offset += 4
            except:
                return False
        stat = inp(ata_reg['alt']['address'])
        if (stat & ERR):
            return False
    return True

def write_data(buffer):
    count = len(buffer) // 512
    if count < 0:
        return False
    offset = 0
    stat = inp(ata_reg['alt']['address'])
    while (stat & DRQ) or (stat & BSY) and ~(stat & ERR):
        if (stat & BSY):
            start = time()
            while (stat & BSY):
                if time() - start > 3:
                    if offset == len(buffer):
                        break
                    else:
                        return False
                stat = inp(ata_reg['alt']['address'])
        if (stat & DRQ):
            try:
                if is_64bit_os():
                    data = struct.unpack_from('@L', buffer, offset)[0]
                    outpd(ata_reg['data']['address'], data)
                    offset += 4
                else:
                    data = struct.unpack_from('@H', buffer, offset)[0]
                    outpw(ata_reg['data']['address'], data)
                    offset += 2
            except:
                return False
        stat = inp(ata_reg['alt']['address'])
        if (stat & ERR):
            return False
    return True

def is_64bit_os():
    if "PROCESSOR_ARCHITEW6432" in os.environ:
        if os.environ["PROCESSOR_ARCHITEW6432"] == "AMD64":
               return True
    if "64" in os.environ["PROCESSOR_ARCHITECTURE"]:
        return True
    return False

def get_identify():
    ident_buffer = bytearray(512)
    if get_identify_data(ident_buffer):
        return dec_identify(ident_buffer)
    return 'x'

def get_identify_data(buffer):
    global ident_dev_data
    if cmd_identify() and read_data(buffer):
        memset(addressof(ident_dev_data), 0, sizeof(ident_dev_data))
        memmove(addressof(ident_dev_data), byref(c_buffer(bytes(buffer))), \
            len(buffer))
        return True
    return False

def get_char(buffer, offset, size):
    temp = ''
    j = 0
    r = size // 2
    for i in range(r):
        temp += chr(buffer[offset + j + 1])
        temp += chr(buffer[offset + j])
        j += 2
    return temp.strip()

def dec_identify(buffer):
    ident = '\nDrive Identify Data'
    ident += '\n------------------------------------------------------------'
    model = get_char(buffer, 54, 40)
    fwrev = get_char(buffer, 46, 8)
    serial = get_char(buffer, 20, 20)
    temp = '\nModel Number   : ' + model
    ident += temp
    temp = '\nFirmware Rev.  : ' + fwrev
    ident += temp
    temp = '\nSerial Number  : ' + serial
    ident += temp
    lba2 = struct.unpack_from('@LL', buffer, 200)
    lba = lba2[0] + (lba2[1] << 32)
    temp = '\nUser Sectors   : '+str(lba)+' ('+str(int(lba*512/1000000000))+'G)'
    uflag = buffer[82] + (buffer[83] << 8)
    if uflag & 0x400:
        temp += ' - HPA'
    ident += temp
    temp = '\nBuffers        : ' + str(int((buffer[42]+(buffer[43]<<8))/2))+'KB'
    ident += temp
    # Security Mode feature set
    uflag =  buffer[164] + (buffer[165] << 8)
    if uflag & 2:
        temp = '\nFeature Sets   : (0x%04X)security supported' % uflag
    else:
        temp = '\nFeature Sets   : (0x%04X)security unsupported' % uflag
    ident += temp
    # Security status
    uflag = buffer[256] + (buffer[257] << 8)
    if uflag & 1:
        temp = '\nSecurity Status: (0x%04X)security supported' % uflag
        if uflag & 2:
            temp += ', enabled'
            if uflag & 256:
                temp += ', level is maximum'
            else:
                temp += ', level is high'
        elif uflag & 4:
            temp += ', locked'
        elif uflag & 8:
            temp += ', frozen'
        elif uflag & 0x10:
            temp += ', count has expired'
        ident += temp
        # Time required for Security erase unit completion
        if uflag & 0x20:
            temp = '\nEnhanced Erase : '
            value = buffer[180] + (buffer[181] << 8)
            if value == 0:
                temp += 'value not specified'
            elif value == 255:
                temp += '>508 minutes'
            else:
                temp = temp + str(value * 2) + ' minutes'
            ident += temp
        temp = '\nNormal Erase   : '
        value = buffer[178] + (buffer[179] << 8)
        if value == 0:
            temp += 'value not specified'
        elif value == 255:
            temp += '>508 minutes'
        else:
            temp = temp + str(value * 2) + ' minutes'
        ident += temp
    # Display Automatic Acoustic Managment
    temp = '\nCurrent AAM Lvl: '
    uflag = buffer[166] + (buffer[167] << 8)
    if uflag & 0x200:
        uflag = buffer[188]
        temp = temp + str(uflag) + ' '
        if uflag & 0xFF:
            temp += '(Reserved)'
        elif uflag & 0xFE:
            temp += '(Performance)'
        elif uflag >= 0x81:
            temp += '(Intermediate)'
        elif uflag == 0x80:
            temp += '(Quiet Mode)'
        elif uflag >= 1:
            temp += '(Retired)'
        else:
            temp += '(Reserved)'
        uflag = buffer[172] + (buffer[173] << 8)
        if uflag & 0x200 == 0:
            temp += ', AAM is disabled'
        else:
            temp += ', AAM is enabled'
    else:
        temp += 'AAM not supported'
    ident += temp
    # Display UDMA Mode
    temp = '\nUdma Mode      : '
    uflag = buffer[176]
    if uflag == 0xFF:
        temp += 'Disabled'
    else:
        for i in range(9):
            if uflag & 1:
                temp = temp + str(i) + ' '
            uflag >>= 1
    ident += temp
    # Display Cache State
    temp = '\nCache State    : '
    uflag = buffer[170]
    if uflag & 64:
        temp += 'RCache On, '
    else:
        temp += 'RCache Off, '
    if uflag & 32:
        temp += 'WCache On'
    else:
        temp += 'WCache Off'
    ident += temp
    # Display Smart State
    temp = '\nSmart Enabled  : '
    if uflag & 1:
        temp += 'On'
    else:
        temp += 'Off'
    ident += temp
    # Create Directory
    return ident

